#include <jni.h>
#include <libyuv.h>
#include "libyuv/ext/ext_convert.h"
#include "libyuv_ext.h"
#include <android/log.h>

#ifdef __cplusplus
extern "C" {
#endif

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_ABGRToARGB(JNIEnv *env, jobject, jbyteArray src_rgba, jbyteArray dst_bgra, jint width, jint height) {
    jbyte *abgrbuf = env->GetByteArrayElements(src_rgba, nullptr);
    jbyte *argbbuf = env->GetByteArrayElements(dst_bgra, nullptr);
    libyuv::ABGRToARGB(
            (const uint8_t *) abgrbuf, width * 4,
            (uint8_t *) argbbuf, width * 4,
            width, height
    );
    env->ReleaseByteArrayElements(src_rgba, abgrbuf, 0);
    env->ReleaseByteArrayElements(dst_bgra, argbbuf, 0);
}

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_BGRToNV21(JNIEnv *env, jobject, jbyteArray src_rgb, jbyteArray dst_yuv, jint width, jint height) {
    jbyte *bgrbuf = env->GetByteArrayElements(src_rgb, nullptr);
    jbyte *yuvbuf = env->GetByteArrayElements(dst_yuv, nullptr);
    int byte_width = width * 3;
    width -= width % 2;
    height -= height % 2;
    int len = width * height;
    auto *des_y = (uint8_t *) yuvbuf;
    uint8_t *des_vu = des_y + len;

    libyuv::BGR24ToNV21(
            (const uint8_t *) bgrbuf, byte_width,
            des_y, width,
            des_vu, width,
            width, height
    );

    env->ReleaseByteArrayElements(src_rgb, bgrbuf, 0);
    env->ReleaseByteArrayElements(dst_yuv, yuvbuf, 0);
}

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_ABGRToNV21(JNIEnv *env, jobject, jbyteArray src_rgba, jbyteArray dst_yuv, jint width, jint height) {
    jbyte *abgrbuf = env->GetByteArrayElements(src_rgba, nullptr);
    jbyte *yuvbuf = env->GetByteArrayElements(dst_yuv, nullptr);
    int byte_width = width * 4;
    width -= width % 2;
    height -= height % 2;
    int len = width * height;
    auto *des_y = (uint8_t *) yuvbuf;
    uint8_t *des_vu = des_y + len;

#if 0
    uint8_t *des_u = des_y + len;
    uint8_t *des_v = des_u + len / 4;
#endif

    /**
     * 需要注意的是 opengles取出来的数据存储序列为：RRRRRRRRGGGGGGGGBBBBBBBBAAAAAAAA
     * 但实际上这种序列算是ABGR的存储序列
     */
    libyuv::ABGRToNV21(
            (const uint8_t *) abgrbuf, byte_width,
            des_y, width,
            des_vu, width,
            width, height
    );

    env->ReleaseByteArrayElements(src_rgba, abgrbuf, 0);
    env->ReleaseByteArrayElements(dst_yuv, yuvbuf, 0);
}

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_ScaleABGRToNV21(JNIEnv *env, jobject, jbyteArray src_rgba, jbyteArray dst_yuv, jint width, jint height, jfloat scale) {
    jbyte *yuvbuf = env->GetByteArrayElements(dst_yuv, nullptr);
    jbyte *abgrbuf = env->GetByteArrayElements(src_rgba, nullptr);

    jbyteArray original_argb = env->NewByteArray(env->GetArrayLength(src_rgba));
    jbyte *original_argbbuf = env->GetByteArrayElements(original_argb, nullptr);
    libyuv::ABGRToARGB((const uint8_t *) abgrbuf, width * 4, (uint8_t *) original_argbbuf, width * 4, width, height);

    int scale_stride = width * 4 * scale;
    int scale_width = width * scale;
    int scale_height = height * scale;
    jbyteArray scale_argb = env->NewByteArray(env->GetArrayLength(src_rgba) * scale);
    jbyte *scale_argbbuf = env->GetByteArrayElements(scale_argb, nullptr);
    libyuv::ARGBScale((uint8_t *) original_argbbuf, width * 4, width, height, (uint8_t *) scale_argbbuf, scale_stride, scale_width, scale_height, libyuv::FilterMode::kFilterNone);

    int scale_len = scale_width * scale_height;
    auto *des_y = (uint8_t *) yuvbuf;
    uint8_t *des_vu = des_y + scale_len;
    libyuv::ARGBToNV21((uint8_t *) scale_argbbuf, scale_stride, des_y, scale_width, des_vu, scale_width, scale_width, scale_height);

    env->ReleaseByteArrayElements(dst_yuv, yuvbuf, 0);
    env->ReleaseByteArrayElements(src_rgba, abgrbuf, 0);
    env->ReleaseByteArrayElements(scale_argb, scale_argbbuf, 0);
    env->ReleaseByteArrayElements(original_argb, original_argbbuf, 0);
}

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_ScaleARGBToNV21(JNIEnv *env, jobject, jbyteArray src_rgba, jbyteArray dst_yuv, jint width, jint height, jfloat scale) {
    jbyte *yuvbuf = env->GetByteArrayElements(dst_yuv, nullptr);
    jbyte *abgrbuf = env->GetByteArrayElements(src_rgba, nullptr);

    int scale_stride = width * 4 * scale;
    int scale_width = width * scale;
    int scale_height = height * scale;
    jbyteArray scale_argb = env->NewByteArray(env->GetArrayLength(src_rgba) * scale);
    jbyte *scale_argbbuf = env->GetByteArrayElements(scale_argb, nullptr);
    libyuv::ARGBScale((uint8_t *) abgrbuf, width * 4, width, height, (uint8_t *) scale_argbbuf, scale_stride, scale_width, scale_height, libyuv::FilterMode::kFilterNone);

    int scale_len = scale_width * scale_height;
    auto *des_y = (uint8_t *) yuvbuf;
    uint8_t *des_vu = des_y + scale_len;
    libyuv::ARGBToNV21((uint8_t *) scale_argbbuf, scale_stride, des_y, scale_width, des_vu, scale_width, scale_width, scale_height);

    env->ReleaseByteArrayElements(dst_yuv, yuvbuf, 0);
    env->ReleaseByteArrayElements(src_rgba, abgrbuf, 0);
    env->ReleaseByteArrayElements(scale_argb, scale_argbbuf, 0);
}

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_ScaleYUV(JNIEnv *env, jobject, jbyteArray src_yuv, jint src_width, jint src_height, jbyteArray dst_yuv, jfloat scale) {
    jbyte *src_yuvbuf = env->GetByteArrayElements(src_yuv, nullptr);
    jbyte *dst_yuvbuf = env->GetByteArrayElements(dst_yuv, nullptr);
    int scale_width = src_width * scale;
    int scale_height = src_height * scale;
#if 0
    jbyteArray dst_yuv = env->NewByteArray(env->GetArrayLength(src_yuv) * scale);
    jbyte *dst_yuvbuf = env->GetByteArrayElements(dst_yuv, nullptr);
#endif
    libyuv::ScalePlane((uint8_t *) src_yuvbuf, src_width, src_width, src_height, (uint8_t *) dst_yuvbuf, scale_width, scale_width, scale_height, libyuv::FilterMode::kFilterNone);
    env->ReleaseByteArrayElements(src_yuv, src_yuvbuf, 0);
    env->ReleaseByteArrayElements(dst_yuv, dst_yuvbuf, 0);
}

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_ScaleABGR(JNIEnv *env, jobject, jbyteArray src_abgr, jint src_width, jint src_height, jbyteArray dst_abgr, jfloat scale) {
    jbyte *src_abgrbuf = env->GetByteArrayElements(src_abgr, nullptr);
    jbyte *dst_abgrbuf = env->GetByteArrayElements(dst_abgr, nullptr);
    int scale_width = src_width * scale;
    int scale_height = src_height * scale;
    libyuv::ARGBScale((uint8_t *) src_abgrbuf, src_width * 4, src_width, src_height, (uint8_t *) dst_abgrbuf, scale_width * 4, scale_width, scale_height, libyuv::FilterMode::kFilterNone);
    env->ReleaseByteArrayElements(src_abgr, src_abgrbuf, 0);
    env->ReleaseByteArrayElements(dst_abgr, dst_abgrbuf, 0);
}

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_YUVToYUV420(JNIEnv *env, jobject, jbyteArray src_y, jbyteArray src_u, jbyteArray src_v, jbyteArray dst_yuv, jint width, jint height) {
    jbyte *ybuf = env->GetByteArrayElements(src_y, nullptr);
    jbyte *ubuf = env->GetByteArrayElements(src_u, nullptr);
    jbyte *vbuf = env->GetByteArrayElements(src_v, nullptr);
    jbyte *yuvbuf = env->GetByteArrayElements(dst_yuv, nullptr);
    width -= width % 2;
    height -= height % 2;

#if 0
    int len = width * height;
    auto *des_y = (uint8_t *) yuvbuf;
    uint8_t *des_vu = des_y + len;
    libyuv::I422ToNV21(
            (const uint8_t *) ybuf, width,
            (const uint8_t *) ubuf, width,
            (const uint8_t *) vbuf, width,
            des_y, width,
            des_vu, width,
            width, height
    );
#endif

    libyuv::ConvertFromI420(
            (const uint8_t *) ybuf, width,
            (const uint8_t *) ubuf, (width + 1) / 2,
            (const uint8_t *) vbuf, (width + 1) / 2,
            (uint8_t *) yuvbuf, width,
            width, height, libyuv::FourCC::FOURCC_NV21
    );

    env->ReleaseByteArrayElements(src_y, ybuf, 0);
    env->ReleaseByteArrayElements(src_u, ubuf, 0);
    env->ReleaseByteArrayElements(src_v, vbuf, 0);
    env->ReleaseByteArrayElements(dst_yuv, yuvbuf, 0);
}

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_YUV420ToNV21(JNIEnv *env, jobject, jobject src_y, jobject src_u, jobject src_v, jbyteArray dst_yuv, jint width, jint height) {
    auto *ybuf = (jbyte *) env->GetDirectBufferAddress(src_y);
    auto *ubuf = (jbyte *) env->GetDirectBufferAddress(src_u);
    auto *vbuf = (jbyte *) env->GetDirectBufferAddress(src_v);
    jbyte *yuvbuf = env->GetByteArrayElements(dst_yuv, nullptr);

    int len = width * height;
    auto *des_y = (uint8_t *) yuvbuf;
    uint8_t *des_vu = des_y + len;
#if 0
    jbyteArray src_rgba = env->NewByteArray(len * 4);
    jbyte *rgbabuf = env->GetByteArrayElements(src_rgba, nullptr);
    libyuv::I420ToABGR(
            (const uint8_t *) ybuf, width,
            (const uint8_t *) ubuf, width / 2,
            (const uint8_t *) vbuf, width / 2,
            (uint8_t *) rgbabuf, width * 4,
            width, height
    );
    libyuv::ABGRToNV21(
            (const uint8_t *) rgbabuf, width * 4,
            des_y, width,
            des_vu, width,
            width, height
    );
    env->ReleaseByteArrayElements(src_rgba, rgbabuf, 0);
#else
    libyuv::I420ToNV21(
            (const uint8_t *) ybuf, width,
            (const uint8_t *) ubuf, SUBSAMPLE(width, 2),
            (const uint8_t *) vbuf, SUBSAMPLE(width, 2),
            des_y, width,
            des_vu, width,
            width, height
    );
#endif
    env->ReleaseByteArrayElements(dst_yuv, yuvbuf, 0);
}

#ifdef __cplusplus
}
#endif